package com.opendeepsearch;

import java.util.*;
import java.util.concurrent.*;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import dev.langchain4j.data.message.*;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.openai.OpenAiChatModel;

import com.opendeepsearch.config.Configuration;
import com.opendeepsearch.state.*;
import com.opendeepsearch.tool.ArxivSearchTool;
import com.opendeepsearch.tool.TavilySearchTool;
import com.opendeepsearch.tool.Tool;

public class OpenDeepResearch {
    private static final Logger logger = LoggerFactory.getLogger(OpenDeepResearch.class);
    
    private final Configuration configuration;
    private final Map<String, ChatModel> models;
    private final Map<String, Tool> tools;
    private final ExecutorService executorService;

    public OpenDeepResearch(Configuration configuration) {
            this.configuration = configuration;
            this.models = initializeModels();
            this.tools = initializeTools();
            this.executorService = Executors.newFixedThreadPool(configuration.getMaxConcurrentResearchUnits());
        }

    // 初始化模型
    private Map<String, ChatModel> initializeModels() {
        Map<String, ChatModel> modelMap = new HashMap<>();
        String apiKey = System.getenv("KIMI_API_KEY");
        if (apiKey == null || apiKey.trim().isEmpty()) {
            throw new IllegalStateException("KIMI_API_KEY environment variable is required but not set");
        }
        
        ChatModel kimiChatModel = OpenAiChatModel.builder()
               .apiKey(apiKey)
               .baseUrl("https://api.moonshot.cn/v1")
               .modelName(configuration.getResearchModel())
               .build();
        
        // 使用配置中的模型名称作为key
        modelMap.put(configuration.getResearchModel(), kimiChatModel);
        modelMap.put(configuration.getSummarizationModel(), kimiChatModel);
        modelMap.put(configuration.getCompressionModel(), kimiChatModel);
        modelMap.put(configuration.getFinalReportModel(), kimiChatModel);
        
        logger.info("Initialized {} models", modelMap.size());
        return modelMap;
    }

    // 初始化工具
    private Map<String, Tool> initializeTools() {
        Map<String, Tool> toolMap = new HashMap<>();
        
        // 添加Tavily搜索工具
        if (configuration.getSearchApi().equalsIgnoreCase("tavily")) {
             String apiKey = System.getenv("TAVILY_API_KEY");
            toolMap.put("tavily_search", new TavilySearchTool(apiKey));
        }
        
        // 添加ArXiv搜索工具
        toolMap.put("arxiv_search", new ArxivSearchTool());
        
        return toolMap;
    }

    // 执行完整的研究流程
    public String executeResearch(String userQuery) throws Exception {
        if (userQuery == null || userQuery.trim().isEmpty()) {
            throw new IllegalArgumentException("User query cannot be null or empty");
        }
        
        logger.info("Starting research for query: {}", userQuery);
        
        try {
            // 1. 处理用户输入
            AgentInputState inputState = new AgentInputState(
                UUID.randomUUID().toString(), "INITIALIZED", userQuery.trim(), Collections.emptyList());
            
            // 2. 澄清问题 (如果配置允许)
            if (configuration.isAllowClarification()) {
                String clarification = clarifyWithUser(inputState);
                logger.info("Clarification needed: {}", clarification);
            }
            
            // 3. 编写研究简报
            SupervisorState supervisorState = writeResearchBrief(inputState);
            logger.info("Research brief generated with {} questions", supervisorState.getResearchQuestions().size());
            
            // 4. 执行研究
            List<ResearcherState> researcherStates = conductResearch(supervisorState);
            
            // 5. 压缩研究结果
            String compressedResults = compressResearchResults(researcherStates);
            
            // 6. 生成最终报告
            String finalReport = generateFinalReport(supervisorState.getResearchBrief(), compressedResults);
            
            logger.info("Research completed successfully");
            return finalReport;
            
        } catch (Exception e) {
            logger.error("Research execution failed", e);
            throw new Exception("Failed to execute research", e);
        }
    }

    private String generateChatResponse(String prompt, String modelKey) {
        ChatModel model = models.get(modelKey);
        if (model == null) {
            throw new IllegalStateException("Model not found: " + modelKey);
        }
        
        List<ChatMessage> messages = Arrays.asList(UserMessage.from(prompt));
        ChatResponse response = model.chat(messages);
        return response.aiMessage().text();
    }

    // 澄清问题
    private String clarifyWithUser(AgentInputState inputState) {
        String prompt = "用户查询: " + inputState.getUserQuery() + 
                      "\n请确认是否需要澄清以下问题: " + inputState.getContext();
        return generateChatResponse(prompt, configuration.getResearchModel());
    }

    // 编写研究简报
    private SupervisorState writeResearchBrief(AgentInputState inputState) {
        String prompt = "将以下用户查询转换为详细的研究简报: " + inputState.getUserQuery() +
                      "\n上下文信息: " + inputState.getContext();
        
        String researchBrief = generateChatResponse(prompt, configuration.getResearchModel());
        
        // 提取研究问题
        String questionsPrompt = "从研究简报中提取关键研究问题:\n" + researchBrief;
        String questionsResponse = generateChatResponse(questionsPrompt, configuration.getResearchModel());
        
        List<String> researchQuestions = Arrays.stream(questionsResponse.split("\n"))
               .filter(s -> !s.trim().isEmpty())
               .map(String::trim)
               .collect(Collectors.toList());
        
        return new SupervisorState(UUID.randomUUID().toString(), "BRIEF_GENERATED", 
                researchBrief, researchQuestions, 0);
    }

    // 执行研究
    private List<ResearcherState> conductResearch(SupervisorState supervisorState) throws Exception {
        List<ResearcherState> researcherStates = new ArrayList<>();
        List<Future<ResearcherState>> futures = new ArrayList<>();
        
        // 为每个研究问题创建一个研究任务
        for (String question : supervisorState.getResearchQuestions()) {
            futures.add(executorService.submit(() -> {
                ResearcherState researcherState = new ResearcherState(
                        UUID.randomUUID().toString(), 
                        "RESEARCHING", 
                        question, 
                        new ArrayList<>(), 
                        ""
                );
                
                try {
                    // 执行多次迭代的研究
                    for (int i = 0; i < configuration.getMaxResearcherIterations(); i++) {
                        // 执行搜索
                        searchAndProcessResults(researcherState);
                        
                        // 检查是否需要进一步研究
                        if (!shouldContinueResearch(researcherState)) {
                            break;
                        }
                    }
                    
                    // 生成研究摘要
                    generateResearchSummary(researcherState);
                    
                    researcherState.setStatus("COMPLETED");
                } catch (Exception e) {
                    researcherState.setStatus("FAILED");
                    e.printStackTrace();
                }
                
                return researcherState;
            }));
        }
        
        // 收集所有研究结果
        for (Future<ResearcherState> future : futures) {
            researcherStates.add(future.get());
        }
        
        return researcherStates;
    }

    // 执行搜索并处理结果
    private void searchAndProcessResults(ResearcherState researcherState) throws Exception {
        Tool searchTool = tools.get("tavily_search");
        if (searchTool == null) {
            throw new IllegalStateException("Tavily search tool not found");
        }
        
        String searchResults = searchTool.execute(researcherState.getResearchQuestion());
        processSearchResults(researcherState, searchResults);
    }

    // 处理搜索结果
    private void processSearchResults(ResearcherState researcherState, String searchResults) {
        try {
            List<SearchResult> results = com.opendeepsearch.util.JsonUtil.parseTavilySearchResults(searchResults);
            researcherState.getSearchResults().addAll(results);
            logger.debug("Added {} search results for question: {}", results.size(), researcherState.getResearchQuestion());
        } catch (Exception e) {
            logger.error("Failed to process search results", e);
        }
    }

    // 检查是否应该继续研究
    private boolean shouldContinueResearch(ResearcherState researcherState) {
        // 简化实现，实际应该基于搜索结果的质量和数量来决定
        return researcherState.getSearchResults().size() < 5;
    }

    // 生成研究摘要
    private void generateResearchSummary(ResearcherState researcherState) {
        StringBuilder contextBuilder = new StringBuilder();
        for (SearchResult result : researcherState.getSearchResults()) {
            contextBuilder.append("标题: ").append(result.getTitle()).append("\n");
            contextBuilder.append("内容: ").append(result.getContent()).append("\n\n");
        }
        
        String prompt = "根据以下研究材料，撰写一个简洁的摘要回答问题: " + 
                      researcherState.getResearchQuestion() + 
                      "\n\n研究材料:\n" + contextBuilder.toString();
        
        String summary = generateChatResponse(prompt, configuration.getSummarizationModel());
        researcherState.setSummary(summary);
        logger.debug("Generated summary for question: {}", researcherState.getResearchQuestion());
    }

    // 压缩研究结果
    private String compressResearchResults(List<ResearcherState> researcherStates) {
        StringBuilder allSummaries = new StringBuilder();
        for (ResearcherState state : researcherStates) {
            allSummaries.append("问题: ").append(state.getResearchQuestion()).append("\n");
            allSummaries.append("摘要: ").append(state.getSummary()).append("\n\n");
        }
        
        String prompt = "压缩以下研究摘要，提取关键信息和见解:\n" + allSummaries.toString();
        return generateChatResponse(prompt, configuration.getCompressionModel());
    }

    // 生成最终报告
    private String generateFinalReport(String researchBrief, String compressedResults) {
        String prompt = "根据研究简报和压缩的研究结果，撰写一份全面的研究报告:\n" +
                      "研究简报:\n" + researchBrief + "\n\n" +
                      "压缩的研究结果:\n" + compressedResults + "\n\n" +
                      "报告应包括引言、主要发现、分析和结论。";
        return generateChatResponse(prompt, configuration.getFinalReportModel());
    }

    // 关闭资源
    public void shutdown() {
        executorService.shutdown();
        try {
            if (!executorService.awaitTermination(800, TimeUnit.MILLISECONDS)) {
                executorService.shutdownNow();
            }
        } catch (InterruptedException e) {
            executorService.shutdownNow();
        }
    }

    public static void main(String[] args) {
        try {
            Configuration config = Configuration.fromEnvironment();
            OpenDeepResearch research = new OpenDeepResearch(config);
            
            String userQuery = args.length > 0 ? args[0] : "人工智能在医疗领域的应用现状和未来趋势";
            logger.info("Starting research for query: {}", userQuery);
            
            String report = research.executeResearch(userQuery);
            System.out.println("最终研究报告:\n" + report);
            
        } catch (Exception e) {
            logger.error("Research failed", e);
            System.err.println("Error: " + e.getMessage());
            System.exit(1);
        }
    }
}